import os
import cv2

# 设置输入和输出目录
video_dir = 'data/starai_datasets/data/starai_datasets/videos_train'
image_dir = 'task/train'

# 如果输出目录不存在，则创建
if not os.path.exists(image_dir):
    os.makedirs(image_dir)

# 遍历视频目录中的所有文件
for video_file in os.listdir(video_dir):
    if video_file.endswith('.mp4'):
        video_path = os.path.join(video_dir, video_file)
        cap = cv2.VideoCapture(video_path)
        frame_count = 1
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 保存帧到输出目录
            frame_filename = os.path.join(image_dir, f"{os.path.splitext(video_file)[0]}_frame_{frame_count:04d}.jpg")
            cv2.imwrite(frame_filename, frame)
            frame_count += 1
        
        cap.release()
        print(f"已转换 {video_file} 为 {frame_count - 1} 帧。")

print("所有视频已处理完成。")